
#include <vector>
#include "../crypt/RSAKeyPairGenerator.h"
#include "../crypt/DSAKeyPairGenerator.h"
#include "../crypt/ECDSAKeyPairGenerator.h"
#include "../crypt/AESSymmetricKey.h"
#include "NewKeyDB.h"

/** Constructs the key database interface, opens the key database and creates
 * four KeyTableView classes, populated by EncryptionKey data, DecryptionKey
 * data, SigningKey data and CheckingKey data.
 *
 * This opens or creates a database of keys. The filename should not have any
 * extension. If the .db is not found, the .sql will be used to populate the
 * database. If the .sql also does not exist, a new database will be created.
 * 
 * This also looks for the keys of main actors in the database, and creates
 * new key pairs if they are not present.
 * 
 * @param filename
 * The name of the files containing the database. As this refers to both .sql
 * and .db files, this should not have a file type extension.
 * 
 * @param WIPpwd
 * The work-in-progress password for keys that are in use. The work-in-progress
 * password is a memory resident password that prevents keys that need to be
 * readily available in memory from being written to disk unencrypted when the
 * program shuts down. Keys in the key databases that need to be secret are
 * encrypted using this password, which must be resupplied in the event that
 * the program shuts down.
 * 
 * @param ownname
 * The main actor is identified by this name, which allows this key to be
 * treated specially as this main key is needed to authenticate the program.
 * 
 * @param ownpwd
 * The main actor has this as the password needed to load its own key.
 * 
 * @param persists
 * A boolean flag that sets whether to save the database to disk. */
NewKeyDB::NewKeyDB(const char *filename, const char *WIPpwd, const char *ownname, const char *ownpwd, int persists) : SQLiteDB(filename, persists)
{
   keytablenames = getTableNamesFromCreateTablesIfNecessary();
   encryptionkeys = new KeyTableView(this, KeyTableView::KEY_TABLE_USE_TYPE__ENCRYPTION,
      ENCRYPTION_KEYTABLE_NAME, WIPpwd, ownname, ownpwd);
   decryptionkeys = new KeyTableView(this, KeyTableView::KEY_TABLE_USE_TYPE__DECRYPTION, 
      DECRYPTION_KEYTABLE_NAME, WIPpwd, ownname, ownpwd);
   signingkeys = new KeyTableView(this, KeyTableView::KEY_TABLE_USE_TYPE__SIGNING, 
      SIGNING_KEYTABLE_NAME, WIPpwd, ownname, ownpwd);
   checkingkeys = new KeyTableView(this, KeyTableView::KEY_TABLE_USE_TYPE__CHECKING, 
      CHECKING_KEYTABLE_NAME, WIPpwd, ownname, ownpwd);
   if ((encryptionkeys->isGeneratingNewOwnKeysAcceptable()) &&
       (decryptionkeys->isGeneratingNewOwnKeysAcceptable()))
   {
      EncryptionKey *enckey = NULL;
      DecryptionKey *deckey = NULL;
      int retval = generateNewEncryptionDecryptionKeys(true, ownname, &enckey, &deckey);
      // Insert these keys into the tables next!
      if (retval == 0)
      {
         encryptionkeys->tryToInsertKeyIntoTable(ownname, ownpwd, KEYMARK_CHAINROOT, NULL, enckey);
         decryptionkeys->tryToInsertKeyIntoTable(ownname, ownpwd, KEYMARK_CHAINROOT, NULL, deckey);
         encryptionkeys->retryGettingOwnKeys(ownpwd);
         decryptionkeys->retryGettingOwnKeys(ownpwd);
      } else
      {
         errprint("failed to generate new encryption/decryption keys for self '%s'", ownname);
      }
   }
   if ((signingkeys->isGeneratingNewOwnKeysAcceptable()) &&
       (checkingkeys->isGeneratingNewOwnKeysAcceptable()))
   {
      SigningKey *sgnkey = NULL;
      CheckingKey *chkkey = NULL;
      int retval = generateNewSigningCheckingKeys(true, ownname, &sgnkey, &chkkey);
      // Insert these keys into the tables next!
      if (retval == 0)
      {
         signingkeys->tryToInsertKeyIntoTable(ownname, ownpwd, KEYMARK_CHAINROOT, NULL, sgnkey);
         checkingkeys->tryToInsertKeyIntoTable(ownname, ownpwd, KEYMARK_CHAINROOT, NULL, chkkey);
         signingkeys->retryGettingOwnKeys(ownpwd);
         checkingkeys->retryGettingOwnKeys(ownpwd);
      } else
      {
         errprint("failed to generate new signing/checking keys for self '%s'", ownname);
      }
   }
}

/** Destroys the key database views and interface. */
NewKeyDB::~NewKeyDB()
{
   if (encryptionkeys) delete encryptionkeys;
   encryptionkeys = NULL;
   if (decryptionkeys) delete decryptionkeys;
   decryptionkeys = NULL;
   if (signingkeys) delete signingkeys;
   signingkeys = NULL;
   if (checkingkeys) delete checkingkeys;
   checkingkeys = NULL;
   // doesn't need any destroy string list because all the string are const
   if (keytablenames) delete keytablenames;
   keytablenames = NULL;
}

/** Generate a key pair for encryption and decryption into the database with
 * a given name.
 * 
 * @param ischained
 * A boolean flag that states whether these keys are to be inserted as part of
 * a key chain. A key chain is a set of keys owned by an actor, which begins
 * with a root key. Keys are added through time and removed when their
 * reference count reaches zero.
 * 
 * @param isexpendable
 * A boolean flag that states whether these keys are expendable. If so, the
 * key will be slowly moved into obsolescence.
 * 
 * @param name
 * The name for the key.
 * 
 * @param pwd
 * A password if required. It can be NULL, to signify no password.
 * 
 * @param owner
 * The actor that will own the key. Having this key to encrypt or sign data
 * proves that the actor in question is legitimately working on behalf of the
 * key owner. NULL signifies that the name of the owner is the same as the
 * name of the key.
 *
 * @return
 * Returns zero when the keys were successfully generated and added to their
 * respective tables. This will return non-zero on an error. */
int NewKeyDB::generateEncDecPairIntoDB(int ischained, int isexpendable,
   const char *name, const char *pwd, const char *owner)
{
   // Chained means limited expendability
   int numberofancestors = encryptionkeys->countKeyChainOfOwner(owner) +
                           decryptionkeys->countKeyChainOfOwner(owner);
   int ischainroot = (numberofancestors == 0) && ischained;
   EncryptionKey *enckey = NULL;
   DecryptionKey *deckey = NULL;
   int retval = generateNewEncryptionDecryptionKeys(ischainroot, name, &enckey, &deckey);
   if (retval == 0)
   {
      char *latest_by_owner_enc = NULL;
      char mark_of_latest_by_owner_enc = '\0';
      char *latest_by_owner_dec = NULL;
      char mark_of_latest_by_owner_dec = '\0';
      if ((ischained) && (!ischainroot))
      {
         // Latest keys in existence must be deprecated when required
         latest_by_owner_enc =
            encryptionkeys->getLatestKeyNameOfOwner(owner);
         mark_of_latest_by_owner_enc =
            encryptionkeys->getMarkOfKey(latest_by_owner_enc);
         latest_by_owner_dec =
            decryptionkeys->getLatestKeyNameOfOwner(owner);
         mark_of_latest_by_owner_dec =
            decryptionkeys->getMarkOfKey(latest_by_owner_dec);
      }
      char mark_of_this = KEYMARK_PERMANENT;
      mark_of_this = (isexpendable) ? KEYMARK_EXPENDABLE : mark_of_this;
      mark_of_this = (ischainroot) ? KEYMARK_CHAINROOT : mark_of_this;
      mark_of_this = (ischained) ? KEYMARK_CHAINLATEST : mark_of_this;
      sqlite3_mutex_enter(dbmutex);
      if (mark_of_latest_by_owner_enc == KEYMARK_CHAINLATEST)
         encryptionkeys->setMarkOfKey_unsafe(latest_by_owner_enc, KEYMARK_CHAIN);
      if (mark_of_latest_by_owner_dec == KEYMARK_CHAINLATEST)
         decryptionkeys->setMarkOfKey_unsafe(latest_by_owner_dec, KEYMARK_CHAIN);
      encryptionkeys->tryToInsertKeyIntoTable_tsafe(name, pwd, mark_of_this, owner, enckey, false);
      decryptionkeys->tryToInsertKeyIntoTable_tsafe(name, pwd, mark_of_this, owner, deckey, false);
      sqlite3_mutex_leave(dbmutex);
      if (latest_by_owner_enc) free(latest_by_owner_enc);
      if (latest_by_owner_dec) free(latest_by_owner_dec);
      return 0;
   } else
   {
      errprint("problem generating encryption/decryption keys");
      return -1;
   }
}

/** Generate a key pair for signing and checking the signatures on data and
 * place it into the database with a given name.
 * 
 * @param ischained
 * A boolean flag that states whether these keys are to be inserted as part of
 * a key chain. A key chain is a set of keys owned by an actor, which begins
 * with a root key. Keys are added through time and removed when their
 * reference count reaches zero.
 * 
 * @param isexpendable
 * A boolean flag that states whether these keys are expendable. If so, the
 * key will be slowly moved into obsolescence.
 * 
 * @param name
 * The name for the key.
 * 
 * @param pwd
 * A password if required. It can be NULL, to signify no password.
 * 
 * @param owner
 * The actor that will own the key. Having this key to encrypt or sign data
 * proves that the actor in question is legitimately working on behalf of the
 * key owner. NULL signifies that the name of the owner is the same as the
 * name of the key.
 *
 * @return
 * Returns zero when the keys were successfully generated and added to their
 * respective tables. This will return non-zero on an error. */
int NewKeyDB::generateSgnChkPairIntoDB(int ischained, int isexpendable,
   const char *name, const char *pwd, const char *owner)
{
   // Chained means limited expendability
   int numberofancestors = signingkeys->countKeyChainOfOwner(owner) +
                           checkingkeys->countKeyChainOfOwner(owner);
   int ischainroot = (numberofancestors == 0) && ischained;
   SigningKey *sgnkey = NULL;
   CheckingKey *chkkey = NULL;
   int retval = generateNewSigningCheckingKeys(ischainroot, name, &sgnkey, &chkkey);
   if (retval == 0)
   {
      char *latest_by_owner_sgn = NULL;
      char mark_of_latest_by_owner_sgn = '\0';
      char *latest_by_owner_chk = NULL;
      char mark_of_latest_by_owner_chk = '\0';
      if ((ischained) && (!ischainroot))
      {
         // Latest keys in existence must be deprecated when required
         latest_by_owner_sgn =
            signingkeys->getLatestKeyNameOfOwner(owner);
         mark_of_latest_by_owner_sgn =
            signingkeys->getMarkOfKey(latest_by_owner_sgn);
         latest_by_owner_chk =
            checkingkeys->getLatestKeyNameOfOwner(owner);
         mark_of_latest_by_owner_chk =
            checkingkeys->getMarkOfKey(latest_by_owner_chk);
      }
      char mark_of_this = KEYMARK_PERMANENT;
      mark_of_this = (isexpendable) ? KEYMARK_EXPENDABLE : mark_of_this;
      mark_of_this = (ischainroot) ? KEYMARK_CHAINROOT : mark_of_this;
      mark_of_this = (ischained) ? KEYMARK_CHAINLATEST : mark_of_this;
      sqlite3_mutex_enter(dbmutex);
      if (mark_of_latest_by_owner_sgn == KEYMARK_CHAINLATEST)
         signingkeys->setMarkOfKey_unsafe(latest_by_owner_sgn, KEYMARK_CHAIN);
      if (mark_of_latest_by_owner_chk == KEYMARK_CHAINLATEST)
         checkingkeys->setMarkOfKey_unsafe(latest_by_owner_chk, KEYMARK_CHAIN);
      signingkeys->tryToInsertKeyIntoTable_tsafe(name, pwd, mark_of_this, owner, sgnkey, false);
      checkingkeys->tryToInsertKeyIntoTable_tsafe(name, pwd, mark_of_this, owner, chkkey, false);
      sqlite3_mutex_leave(dbmutex);
      if (latest_by_owner_sgn) free(latest_by_owner_sgn);
      if (latest_by_owner_chk) free(latest_by_owner_chk);
      return 0;
   } else
   {
      errprint("problem generating signing/checking keys");
      return -1;
   }
}

/** Creates the key tables if they do not already exist. PLEASE FIX HACK BROKEN TESTING MESS MARKING-WISE!
 * 
 * @return
 * Returns a string list of constant strings which are the names of the
 * database tables. The list, but not the list items, must be deleted
 * by the caller. */
std::vector<const char *> *NewKeyDB::getTableNamesFromCreateTablesIfNecessary()
{
   static const char keytabledefinition[] =
      "(keyname char(64) primary key not null, "
       "keytype char(32), "
       "keydata char(4096), "
       "owner char(64), "
       "keymark char(1), "
       "expirymark char(1), "
       "usedrefcount decimal(10,0), "
       "existsrefcount decimal(10, 0), "
       "rfc822date char(32))";
   std::vector<const char *> *keytables = new std::vector<const char *>();
   keytables->push_back(ENCRYPTION_KEYTABLE_NAME);
   keytables->push_back(DECRYPTION_KEYTABLE_NAME);
   keytables->push_back(SIGNING_KEYTABLE_NAME);
   keytables->push_back(CHECKING_KEYTABLE_NAME);
   keytables->push_back(NULL);
   ensureTablesAreCreated(4,
      keytables->at(0), keytabledefinition,
      keytables->at(1), keytabledefinition,
      keytables->at(2), keytabledefinition,
      keytables->at(3), keytabledefinition);
   return keytables;
}

/** Generate a key pair for encryption and decryption with a given name.
 * 
 * @param ischainroot
 * A boolean flag that states whether these keys are to be part of a chain.
 * This is to inform automated policies on what type of keys to generate.
 * 
 * @param name
 * The name for the key.
 * 
 * @param[out] enckey
 * A pointer to the place to put the pointer to the EncryptionKey object.
 * 
 * @param[out] deckey
 * A pointer to the place to put the pointer to the DecryptionKey object.
 * 
 * @return
 * Returns zero when the keys were successfully generated and added to their
 * respective tables. This will return non-zero on an error. */
int NewKeyDB::generateNewEncryptionDecryptionKeys(int ischainroot, const char *name, EncryptionKey **enckey, DecryptionKey **deckey)
{
   int t_generatefunc =
      (ischainroot) ?
         KEY_GENERATE_TYPE_ROOT_ENCRYPT_DECRYPT :
         KEY_GENERATE_TYPE_LEAF_ENCRYPT_DECRYPT;
   switch (t_generatefunc)
   {
      case Key::KEY_ASYMMETRIC_PAIR_RSA:
      {
         RSAKeyPairGenerator *nrsakp = new RSAKeyPairGenerator(RSA_BITS, name);
         *enckey = (EncryptionKey *)(nrsakp->popPublicKey());
         *deckey = (DecryptionKey *)(nrsakp->popPrivateKey());
         // RSA key pair does not delete
         delete nrsakp;
         break;
      }
      case Key::KEY_SYMMETRIC_AES:
      {
         AESSymmetricKey *naessym = new AESSymmetricKey(AES_BITS, name);
         AESSymmetricKey *naessym_copy = new AESSymmetricKey(*naessym);
         *enckey = (EncryptionKey *)(naessym);
         *deckey = (DecryptionKey *)(naessym_copy);
         break;
      }
      default:
      {
         *enckey = NULL;
         *deckey = NULL;
         break;
      }
   }
   if ((*enckey == NULL) || (*deckey == NULL))
   {
      if (*enckey) delete *enckey;
      if (*deckey) delete *deckey;
      return -1;
   }
   return 0;
}

/** Generate a key pair with a given name for signing and checking
 * signatures on data.
 * 
 * @param ischainroot
 * A boolean flag that states whether these keys are to be part of a chain.
 * This is to inform automated policies on what type of keys to generate.
 * 
 * @param name
 * The name for the key.
 * 
 * @param[out] sgnkey
 * A pointer to the place to put the pointer to the SigningKey object.
 * 
 * @param[out] chkkey
 * A pointer to the place to put the pointer to the CheckingKey object.
 * 
 * @return
 * Returns zero when the keys were successfully generated and added to their
 * respective tables. This will return non-zero on an error. */
int NewKeyDB::generateNewSigningCheckingKeys(int ischainroot, const char *name, SigningKey **sgnkey, CheckingKey **chkkey)
{
   int t_generatefunc =
      (ischainroot) ?
         KEY_GENERATE_TYPE_ROOT_CHECKING_SIGNING :
         KEY_GENERATE_TYPE_LEAF_CHECKING_SIGNING;
   switch (t_generatefunc)
   {
      case Key::KEY_ASYMMETRIC_PAIR_RSA:
      {
         RSAKeyPairGenerator *nrsakp = new RSAKeyPairGenerator(RSA_BITS, name);
         *chkkey = (CheckingKey *)(nrsakp->popPublicKey());
         *sgnkey = (SigningKey *)(nrsakp->popPrivateKey());
         // RSA key pair does not delete
         delete nrsakp;
         break;
      }
      case Key::KEY_ASYMMETRIC_PAIR_DSA:
      {
         DSAKeyPairGenerator *ndsakp = new DSAKeyPairGenerator(DSA_BITS, name);
         *chkkey = (CheckingKey *)(ndsakp->popPublicKey());
         *sgnkey = (SigningKey *)(ndsakp->popPrivateKey());
         // DSA key pair does not delete
         delete ndsakp;
         break;
      }
      case Key::KEY_ASYMMETRIC_PAIR_ECDSA:
      {
         ECDSAKeyPairGenerator *necdsakp = new ECDSAKeyPairGenerator(ECDSA_BITS, name);
         *chkkey = (CheckingKey *)(necdsakp->popPublicKey());
         *sgnkey = (SigningKey *)(necdsakp->popPrivateKey());
         // ECDSA key pair does not delete
         delete necdsakp;
         break;
      }
      default:
      {
         *chkkey = NULL;
         *sgnkey = NULL;
         break;
      }
   }
   if ((*chkkey == NULL) || (*sgnkey == NULL))
   {
      if (*chkkey) delete *chkkey;
      if (*sgnkey) delete *sgnkey;
      return -1;
   }
   return 0;
}
